# app/rag/embedding/openai.py

from .base import EmbeddingProvider
from langchain_openai import OpenAIEmbeddings
from typing import List
from ...core.logger import logger
import os


class OpenAIEmbedding(EmbeddingProvider):
    def __init__(
        self,
        model: str = "text-embedding-3-small",
        api_key: str = None,
        api_base: str = None,
        organization: str = None
    ):
        """
        初始化 OpenAI Embedding
        :param model: 模型名，默认 text-embedding-3-small（性价比高）
        :param api_key: OpenAI API Key（优先从参数读，其次环境变量）
        :param api_base: 自定义 endpoint（如 Azure）
        :param organization: OpenAI 组织 ID（可选）
        """
        api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OPENAI_API_KEY 未提供，请设置环境变量或传参")

        self.model = model
        self.embedding = OpenAIEmbeddings(
            model=model,
            api_key=api_key,
            base_url=api_base,  # 支持自定义 endpoint（如反向代理）
            organization=organization,
            # 请求超时等可扩展参数
            timeout=30,
        )
        logger.info(f"🌐 OpenAI Embedding loaded: {model}")

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        # 过滤空文本
        texts = [t for t in texts if t.strip()]
        if not texts:
            return []
        return self.embedding.embed_documents(texts)

    def embed_query(self, text: str) -> List[float]:
        if not text.strip():
            return [0.0] * 1536  # 假设维度为 1536（实际应动态获取）
        return self.embedding.embed_query(text)